import argparse
import json
import sys
import time
from pathlib import Path
import threading

import numpy as np
import cv2

# 兼容两种常见 Python OpenNI2 封装：primesense（推荐）与 openni
openni_import_error = None
try:
    from primesense import openni2, _openni2 as c_api  # type: ignore
except Exception as exc_prime:
    try:
        from openni import openni2, _ext as c_api  # type: ignore
    except Exception as exc_openni:
        openni_import_error = (exc_prime, exc_openni)
        print("需要 Python OpenNI2 封装 (pip install primesense) 且本机已安装 OpenNI2/Orbbec SDK", file=sys.stderr)
        raise


def try_initialize_openni() -> None:
    # 优先默认搜索路径，其次尝试常见安装目录
    try:
        openni2.initialize()
        return
    except Exception:
        pass
    candidates = [
        r"C:\\Program Files\\OpenNI2\\Redist",
        r"C:\\Program Files (x86)\\OpenNI2\\Redist",
        r"C:\\Program Files\\Orbbec\\OpenNI2\\Redist",
        r"C:\\Program Files (x86)\\Orbbec\\OpenNI2\\Redist",
    ]
    for p in candidates:
        try:
            openni2.initialize(p)
            return
        except Exception:
            continue
    raise RuntimeError("无法初始化 OpenNI2，请确认已安装 OpenNI2/Orbbec SDK 并将 Redist 目录加入 PATH")


def ensure_dir(path: Path) -> None:
    path.mkdir(parents=True, exist_ok=True)


def write_orbslam3_yaml(yaml_path: Path,
                        width: int,
                        height: int,
                        fps: int,
                        fx: float,
                        fy: float,
                        cx: float,
                        cy: float,
                        depth_map_factor: float = 1000.0) -> None:
    content = f"""
# Auto-generated for Orbbec (OpenNI2) RGB-D recording
Camera.type: "PinHole"

Camera.width: {width}
Camera.height: {height}

Camera.fx: {fx:.8f}
Camera.fy: {fy:.8f}
Camera.cx: {cx:.8f}
Camera.cy: {cy:.8f}

Camera.k1: 0.0
Camera.k2: 0.0
Camera.p1: 0.0
Camera.p2: 0.0
Camera.k3: 0.0

Camera.fps: {fps}
Camera.RGB: 1

DepthMapFactor: {depth_map_factor:.1f}

ORBextractor.nFeatures: 1200
ORBextractor.scaleFactor: 1.2
ORBextractor.nLevels: 8
ORBextractor.iniThFAST: 20
ORBextractor.minThFAST: 7
"""
    yaml_path.write_text(content.strip() + "\n", encoding="utf-8")


def record_orbbec_openni(output_dir: Path,
                         width: int,
                         height: int,
                         fps: int,
                         duration: float,
                         fx: float,
                         fy: float,
                         cx: float,
                         cy: float,
                         device_index: int = 0,
                         uvc_index: int = -1) -> None:
    ensure_dir(output_dir)
    rgb_dir = output_dir / "rgb"
    depth_dir = output_dir / "depth"
    ensure_dir(rgb_dir)
    ensure_dir(depth_dir)

    rgb_list = open(output_dir / "rgb.txt", "w", encoding="utf-8")
    depth_list = open(output_dir / "depth.txt", "w", encoding="utf-8")
    assoc_list = open(output_dir / "associations.txt", "w", encoding="utf-8")

    try_initialize_openni()

    # primesense API 常见方式：直接 open_any()，或列 URI（新版本可用 enumerate_uris）
    dev = None
    try:
        dev = openni2.Device.open_any()
    except Exception:
        try:
            # 兼容旧 API：列举 URI 并按索引打开
            enumerate_uris = getattr(openni2.Device, 'enumerate_uris', None)
            if enumerate_uris is not None:
                uris = enumerate_uris()
                if not uris:
                    raise RuntimeError("未检测到 Orbbec/OpenNI2 设备")
                dev = openni2.Device.open(uris[device_index])
            else:
                raise
        except Exception:
            raise RuntimeError("未检测到 Orbbec/OpenNI2 设备，或 OpenNI2 运行库不可用")

    # 检测设备型号（Astra Pro 的 OpenNI2 彩色流会卡死）
    force_uvc = False
    try:
        dev_info = dev.get_device_info()
        usb_pid = dev_info.usbProductId
        # Astra Pro (PID 1027) 的 OpenNI2 彩色流会在 start() 时卡死
        if usb_pid == 1027:
            print(f"检测到 Astra Pro (PID={usb_pid})，OpenNI2 彩色流已知问题，强制使用 UVC")
            force_uvc = True
    except Exception:
        pass

    # 开启对齐（若硬件支持）
    try:
        dev.set_image_registration_mode(openni2.IMAGE_REGISTRATION_DEPTH_TO_COLOR)
    except Exception:
        pass

    depth_stream = dev.create_depth_stream()
    color_stream = None
    if not force_uvc:
        try:
            color_stream = dev.create_color_stream()
        except Exception:
            color_stream = None

    # 配置视频模式
    dvm = c_api.OniVideoMode()
    dvm.pixelFormat = c_api.OniPixelFormat.ONI_PIXEL_FORMAT_DEPTH_1_MM
    dvm.resolutionX = width
    dvm.resolutionY = height
    dvm.fps = fps
    depth_stream.set_video_mode(dvm)

    color_pixel_format = None
    if color_stream is not None:
        # 优先 RGB888，不行则回退 YUYV
        cvm = c_api.OniVideoMode()
        cvm.resolutionX = width
        cvm.resolutionY = height
        cvm.fps = fps
        try:
            cvm.pixelFormat = c_api.OniPixelFormat.ONI_PIXEL_FORMAT_RGB888
            color_stream.set_video_mode(cvm)
            color_pixel_format = 'RGB888'
        except Exception:
            try:
                cvm.pixelFormat = c_api.OniPixelFormat.ONI_PIXEL_FORMAT_YUYV
                color_stream.set_video_mode(cvm)
                color_pixel_format = 'YUYV'
            except Exception:
                pass

    depth_stream.set_mirroring_enabled(False)
    if color_stream is not None:
        color_stream.set_mirroring_enabled(False)

    depth_stream.start()
    
    # 如果强制 UVC，跳过 OpenNI2 彩色流启动
    color_start_success = False
    if force_uvc:
        color_stream = None
        color_pixel_format = None
        print("已跳过 OpenNI2 彩色流，将使用 UVC")
    elif color_stream is not None and color_pixel_format is not None:
        # 正常启动彩色流
        try:
            color_stream.start()
            test_frame = color_stream.read_frame()
            if test_frame:
                color_start_success = True
                print(f"OpenNI2 彩色流启动成功 ({color_pixel_format})")
        except Exception as e:
            print(f"OpenNI2 彩色流启动失败: {e}，将回退到 UVC")
            color_stream = None
            color_pixel_format = None

    depth_scale_factor = 1000.0  # OpenNI2 深度为毫米

    # 保存内参 JSON 与 ORB-SLAM3 YAML
    intr_json = {
        "width": width,
        "height": height,
        "fps": fps,
        "fx": fx,
        "fy": fy,
        "cx": cx,
        "cy": cy,
        "depth_map_factor": depth_scale_factor,
        "depth_unit": "millimeter",
        "model": "Pinhole"
    }
    (output_dir / "camera_intrinsics.json").write_text(json.dumps(intr_json, indent=2), encoding="utf-8")
    write_orbslam3_yaml(output_dir / "orbslam3_orbbec_rgbd.yaml", width, height, fps, fx, fy, cx, cy, depth_map_factor=depth_scale_factor)

    print("开始录制（Q 结束）...")
    start_t = time.time()
    frame_idx = 0
    # 如果没有可用的 OpenNI 彩色流，回退使用 OpenCV UVC 摄像头（自动扫描索引）
    cap = None
    if color_stream is None or not color_start_success:
        def try_open(index: int) -> cv2.VideoCapture | None:
            cam = cv2.VideoCapture(index, cv2.CAP_DSHOW)
            if not cam or not cam.isOpened():
                try:
                    if cam:
                        cam.release()
                except Exception:
                    pass
                return None
            cam.set(cv2.CAP_PROP_FRAME_WIDTH, width)
            cam.set(cv2.CAP_PROP_FRAME_HEIGHT, height)
            cam.set(cv2.CAP_PROP_FPS, fps)
            return cam

        if uvc_index >= 0:
            cap = try_open(uvc_index)
        else:
            for idx in range(0, 6):
                cap = try_open(idx)
                if cap is not None:
                    break

    try:
        while True:
            dframe = depth_stream.read_frame()
            # 读取彩色帧
            if color_stream is not None:
                cframe = color_stream.read_frame()
                cdata = np.frombuffer(cframe.get_buffer_as_uint8(), dtype=np.uint8)
                if color_pixel_format == 'RGB888':
                    color = cdata.reshape((height, width, 3))
                    color_bgr = cv2.cvtColor(color, cv2.COLOR_RGB2BGR)
                elif color_pixel_format == 'YUYV':
                    color = cdata.reshape((height, width, 2))
                    color_bgr = cv2.cvtColor(color, cv2.COLOR_YUV2BGR_YUY2)
                else:
                    # 未知格式，跳过该帧
                    continue
            else:
                if cap is None or not cap.isOpened():
                    raise RuntimeError("未找到可用的彩色视频源（OpenNI2 Color/ UVC）")
                ret, color_bgr = cap.read()
                if not ret:
                    continue
                if color_bgr.shape[1] != width or color_bgr.shape[0] != height:
                    color_bgr = cv2.resize(color_bgr, (width, height))
            if not dframe:
                continue

            ts = time.time()  # OpenNI2 时间戳不统一，这里采用系统时间
            ts_str = f"{ts:.6f}"

            # 深度帧：uint16 毫米
            ddata = np.frombuffer(dframe.get_buffer_as_uint16(), dtype=np.uint16)
            depth_mm = ddata.reshape((height, width))

            # 保存
            rgb_filename = f"{ts_str}.png"
            depth_filename = f"{ts_str}.png"
            cv2.imwrite(str(rgb_dir / rgb_filename), color_bgr)
            cv2.imwrite(str(depth_dir / depth_filename), depth_mm)

            rgb_rel = f"rgb/{rgb_filename}"
            depth_rel = f"depth/{depth_filename}"
            rgb_list.write(f"{ts_str} {rgb_rel}\n")
            depth_list.write(f"{ts_str} {depth_rel}\n")
            assoc_list.write(f"{ts_str} {rgb_rel} {ts_str} {depth_rel}\n")

            frame_idx += 1

            preview = color_bgr.copy()
            cv2.putText(preview, f"{frame_idx}  {ts_str}s", (12, 28), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
            cv2.imshow("Orbbec OpenNI2 Recording (Q to stop)", preview)
            key = cv2.waitKey(1) & 0xFF
            if key in (ord('q'), ord('Q')):
                break

            if duration is not None and duration > 0 and (time.time() - start_t) >= duration:
                break
    finally:
        try:
            if color_stream is not None:
                color_stream.stop()
            depth_stream.stop()
        except Exception:
            pass
        openni2.unload()
        if cap is not None:
            try:
                cap.release()
            except Exception:
                pass
        rgb_list.close()
        depth_list.close()
        assoc_list.close()
        cv2.destroyAllWindows()

    print(f"录制完成，共保存 {frame_idx} 帧。数据集位于: {output_dir}")


def main() -> None:
    parser = argparse.ArgumentParser(description="录制 Orbbec(OpenNI2) RGB-D 为 TUM 数据集，并生成 ORB-SLAM3 配置")
    parser.add_argument("--out", required=True, help="输出数据集目录，例如 data/room1")
    parser.add_argument("--width", type=int, default=640)
    parser.add_argument("--height", type=int, default=480)
    parser.add_argument("--fps", type=int, default=30)
    parser.add_argument("--duration", type=float, default=None, help="录制时长（秒），为空则按 Q 手动停止")
    parser.add_argument("--device_index", type=int, default=0, help="设备索引，多设备时可选")
    parser.add_argument("--uvc_index", type=int, default=-1, help="UVC 彩色视频索引（-1 自动扫描）")
    # 内参（如不传，将使用 configs/rgbd_custom.yaml 的默认值）
    parser.add_argument("--fx", type=float, default=494.17248433)
    parser.add_argument("--fy", type=float, default=494.94792223)
    parser.add_argument("--cx", type=float, default=317.00285128)
    parser.add_argument("--cy", type=float, default=241.76697033)

    args = parser.parse_args()
    out_dir = Path(args.out)
    record_orbbec_openni(out_dir, args.width, args.height, args.fps, args.duration or 0.0,
                         args.fx, args.fy, args.cx, args.cy, device_index=args.device_index,
                         uvc_index=args.uvc_index)


if __name__ == "__main__":
    main()


